#!/usr/bin/env python

from datasource import *
import pyfits
import os, time

class FITSDataSource(DataSource):
	"""A FITS based DataSource, i.e. the content of 1 HDU. The meta data are the FITS headers
	For multi-extension FITS file containing more than one array
	serveral DataSource must be created, one per HDU"""
	def __init__(self, fileName=None, hduNumber=0):
		stats = os.stat(fileName)
		self.lastmodDate = time.localtime(stats[8])
		DataSource(self.lastmodDate)
        # The DAP file URL
		assert fileName!=None
		self.fileName = fileName
		self.hduNumber = hduNumber
		# Open file
		hudList = pyfits.open(self.fileName)
		self.hdu = hudList[hduNumber]

	def __getstate__(self):
		"""Used to Pickle the object"""
		odict = self.__dict__.copy()	# copy the dict since we change it
		del odict['hdu']				# remove hdu entry
		return odict

	def __setstate__(self,dict):
		"""Used to UnPickle the object"""
		hudList = pyfits.open(dict['fileName'])
		self.__dict__.update(dict)   # update attributes
		self.hdu = hudList[self.hduNumber]
        
	def getMetaData(self):
		"""Return the key:value dictionary containing all the metadata associated with this DataSource."""
		return self.hdu.header
		
	def getData(self, slice=None):
		"""Return the numpy array of this DataSource or a slice of it"""
		if slice==None:
			return self.hdu.data
		else:
			return self.hdu.data[slice]

	def getDataType(self):
		"""Return the numpy array type for this DataResource"""
		return self.hdu.data.type()
		
	def getDataShape(self):
		"""Return the numpy array shape for this DataResource"""
		return self.hdu.data.shape
	

if __name__ == '__main__':
    myFile = FITSDataSource('data/image.fits')
    print myFile.getDataShape()
    print myFile.getDataType()
    print myFile.getData(slice(1,10))
    print myFile.getMetaData()
    print myFile.lastmodDate